// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/substrait/substrait_extension_collector.h>
#include <pollux/testing/gtest_utils.h>
#include <pollux/core/plan_node.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>

using namespace kumo::pollux;
using namespace kumo::pollux::substrait;

namespace kumo::pollux::substrait::test {
    class SubstraitExtensionCollectorTest : public ::testing::Test {
    protected:
        void SetUp() override {
            Test::SetUp();
            functions::prestosql::registerAllScalarFunctions();
        }

        int getReferenceNumber(
            const std::string &functionName,
            std::vector<TypePtr> &&arguments) {
            int referenceNumber1 =
                    extensionCollector_->getReferenceNumber(functionName, arguments);
            // Repeat the call to make sure properly de-duplicated.
            int referenceNumber2 =
                    extensionCollector_->getReferenceNumber(functionName, arguments);
            EXPECT_EQ(referenceNumber1, referenceNumber2);
            return referenceNumber1;
        }

        int getReferenceNumber(
            const std::string &functionName,
            std::vector<TypePtr> &&arguments,
            core::AggregationNode::Step step) {
            int referenceNumber1 =
                    extensionCollector_->getReferenceNumber(functionName, arguments, step);
            // Repeat the call to make sure properly de-duplicated.
            int referenceNumber2 =
                    extensionCollector_->getReferenceNumber(functionName, arguments, step);
            EXPECT_EQ(referenceNumber1, referenceNumber2);
            return referenceNumber2;
        }

        /// Given a substrait plan, return the sorted extension functions by the
        /// function anchor.
        ::google::protobuf::RepeatedPtrField<
            ::substrait::extensions::SimpleExtensionDeclaration>
        getSortedSubstraitExtension(const ::substrait::Plan *substraitPlan) {
            auto substraitExtensions = substraitPlan->extensions();
            std::sort(
                substraitExtensions.begin(),
                substraitExtensions.end(),
                [](const auto &a, const auto &b) {
                    return a.extension_function().function_anchor() <
                           b.extension_function().function_anchor();
                });

            return substraitExtensions;
        }

        SubstraitExtensionCollectorPtr extensionCollector_ =
                std::make_shared<SubstraitExtensionCollector>();
    };

    TEST_F(SubstraitExtensionCollectorTest, getReferenceNumberForScalarFunction) {
        ASSERT_EQ(getReferenceNumber("plus", {INTEGER(), INTEGER()}), 0);
        ASSERT_EQ(getReferenceNumber("divide", {INTEGER(), INTEGER()}), 1);
        ASSERT_EQ(getReferenceNumber("cardinality", {ARRAY(INTEGER())}), 2);
        ASSERT_EQ(getReferenceNumber("array_sum", {ARRAY(INTEGER())}), 3);

        auto functionType = std::make_shared<const FunctionType>(
            std::vector<TypePtr>{INTEGER(), VARCHAR()}, BIGINT());
        std::vector<TypePtr> types = {MAP(INTEGER(), VARCHAR()), functionType};
        ASSERT_ANY_THROW(getReferenceNumber("transform_keys", std::move(types)));
    }

    TEST_F(
        SubstraitExtensionCollectorTest,
        getReferenceNumberForAggregateFunction) {
        // Sum aggregate function have same argument type for each aggregation step.
        ASSERT_EQ(
            getReferenceNumber(
                "sum", {INTEGER()}, core::AggregationNode::Step::kSingle),
            0);

        // Partial avg aggregate function should use primitive integral type.
        ASSERT_EQ(
            getReferenceNumber(
                "avg", {INTEGER()}, core::AggregationNode::Step::kPartial),
            1);

        // Final avg aggregate function should use struct type, like
        // 'ROW<DOUBLE,BIGINT>'
        ASSERT_EQ(
            getReferenceNumber(
                "avg",
                {ROW({DOUBLE(), BIGINT()})},
                core::AggregationNode::Step::kFinal),
            2);

        // Count aggregate function have same argument type for each aggregation step.
        ASSERT_EQ(
            getReferenceNumber(
                "count", {INTEGER()}, core::AggregationNode::Step::kFinal),
            3);
    }

    TEST_F(SubstraitExtensionCollectorTest, addExtensionsToPlan) {
        getReferenceNumber("plus", {INTEGER(), INTEGER()});
        getReferenceNumber("divide", {INTEGER(), INTEGER()});
        getReferenceNumber("cardinality", {ARRAY(INTEGER())});
        getReferenceNumber("array_sum", {ARRAY(INTEGER())});
        getReferenceNumber("sum", {INTEGER()});
        getReferenceNumber("avg", {INTEGER()});
        getReferenceNumber("avg", {ROW({DOUBLE(), BIGINT()})});
        getReferenceNumber("count", {INTEGER()});

        google::protobuf::Arena arena;
        auto *substraitPlan =
                google::protobuf::Arena::CreateMessage<::substrait::Plan>(&arena);

        extensionCollector_->addExtensionsToPlan(substraitPlan);

        const auto &substraitExtensions = getSortedSubstraitExtension(substraitPlan);
        auto getFunctionName = [&](auto id) {
            return substraitExtensions[id].extension_function().name();
        };

        ASSERT_EQ(substraitPlan->extensions().size(), 8);
        ASSERT_EQ(getFunctionName(0), "plus:i32_i32");
        ASSERT_EQ(getFunctionName(1), "divide:i32_i32");
        ASSERT_EQ(getFunctionName(2), "cardinality:list");
        ASSERT_EQ(getFunctionName(3), "array_sum:list");
        ASSERT_EQ(getFunctionName(4), "sum:i32");
        ASSERT_EQ(getFunctionName(5), "avg:i32");
        ASSERT_EQ(getFunctionName(6), "avg:struct");
        ASSERT_EQ(getFunctionName(7), "count:i32");
    }
} // namespace kumo::pollux::substrait::test
